@echo off
setlocal

echo "python 3.9 version is recommended."

:: 检查 CUDA 版本
where nvcc >nul
if %errorlevel% neq 0 (
    echo CUDA not found. Installing CPU version of PyTorch.
    set PYTORCH_INSTALL_COMMAND=pip3 install torch torchvision torchaudio
) else (
    :: 获取 CUDA 版本
    for /f "tokens=2 delims=." %%i in ('nvcc --version ^| findstr "release"') do set CUDA_VERSION=%%i

    if "%CUDA_VERSION%"=="11.8" (
        set PYTORCH_INSTALL_COMMAND=pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
    ) else if "%CUDA_VERSION%"=="12.1" (
        set PYTORCH_INSTALL_COMMAND=pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121
    ) else if "%CUDA_VERSION%"=="12.4" (
        set PYTORCH_INSTALL_COMMAND=pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu124
    ) else (
        echo Unsupported CUDA version: %CUDA_VERSION%. Installing CPU version of PyTorch.
        set PYTORCH_INSTALL_COMMAND=pip3 install torch torchvision torchaudio
    )
)

:: 安装 PyTorch
echo Installing PyTorch...
%PYTORCH_INSTALL_COMMAND%

:: 安装其他依赖库
echo Installing other dependencies...
pip3 install openai opencv-python pyttsx3 scikit-image ultralytics

echo Setup complete.
endlocal